import os
from pathlib import Path
from typing import List
from docx import Document as DocxDocument
from langchain_core.documents import Document as LangDocument

import os
from typing import List
from docx import Document as DocxDocument
from docx.text.paragraph import Paragraph
from docx.table import Table
from langchain_core.documents import Document as LangDocument

def load_document_with_format(file_path: str) -> List[LangDocument]:
    """加载 Word 文档，保留原格式，包括标题、段落、表格"""
    print(f"加载文档: {file_path}")
    doc = DocxDocument(file_path)
    lines = []

    for block in iter_block_items(doc):
        if isinstance(block, Paragraph):
            text = format_paragraph(block)
            if text:
                lines.append(text)
        elif isinstance(block, Table):
            lines.append("[表格]")
            lines.append(format_table(block))
            lines.append("[/表格]")

    full_text = "\n".join(lines)
    metadata = {
        "source": file_path,
        "file_name": os.path.basename(file_path),
        "file_type": "docx"
    }
    return [LangDocument(page_content=full_text, metadata=metadata)]

# 提取段落 + 表格，保持顺序
from docx.document import Document as _Document

def iter_block_items(parent):
    """按顺序返回段落和表格（混排）"""
    parent_elm = parent.element.body
    for child in parent_elm.iterchildren():
        if child.tag.endswith('}p'):
            yield Paragraph(child, parent)
        elif child.tag.endswith('}tbl'):
            yield Table(child, parent)

def format_paragraph(paragraph: Paragraph) -> str:
    text = paragraph.text.strip()
    if not text:
        return ""
    style = paragraph.style.name
    if style.startswith("Heading"):
        level = style.replace("Heading ", "")
        return f"\n{'#' * int(level)} {text}"
    else:
        return text

def format_table(table: Table) -> str:
    """将表格转换为 Markdown 表格样式"""
    rows_text = []
    for row in table.rows:
        cells = [cell.text.strip().replace("\n", " ") for cell in row.cells]
        rows_text.append(" | ".join(cells))
    return "\n".join(rows_text)


def main():
    path=Path("D:\\hbyt\\AI智能投标\\能力组件-附件\\研发特有_沟通管理及问题事件响应机制_v1.0_2201.docx")
    doc=load_document_with_format(path)
    print(doc)
    print(doc[0].page_content)
if __name__ == '__main__':
    main()